import numpy as np
import os
import cv2 as cv
from shapely.geometry import Polygon, MultiPoint  # 多边形


def compute_one_ration_iou(x_box, y_box):
    # 四边形的二维坐标表示
    xx_box, yy_box = np.array(x_box).reshape(4, 2), np.array(y_box).reshape(4, 2)
    # 构建四边形对象，会自动计算四个点的顺序：左上 左下  右下 右上 左上（返回5个点，最后回到起始点）
    x_poly, y_poly = Polygon(xx_box).convex_hull, Polygon(yy_box).convex_hull

    intersect_area = x_poly.intersection(y_poly).area  # 相交面积
    if intersect_area == 0:
        iou = 0
    else:
        union_area = x_poly.area + y_poly.area - intersect_area  # 总共面积
        iou = intersect_area / union_area
    return iou


def xywh2xy4(box_in, i_size):
    angle = int(box_in[5])
    cx = i_size * float(box_in[1])
    cy = i_size * float(box_in[2])
    w = i_size * float(box_in[3])
    h = i_size * float(box_in[4])
    arg = ((cx, cy), (w, h), angle)
    box = cv.boxPoints(arg)
    return np.int0(box)


def pred_lines2dict(lines_f):
    # 预测结果的字典
    result_dict = {}
    for lll in lines_f:
        lll = lll.split(' ')
        lll[-11] = lll[-11].split('/')[-1]
        lll[-1] = lll[-1][:-1]
        lll[-8:] = [eval(kk) for kk in lll[-8:]]
        lll.pop(-9)
        if lll[-10] not in result_dict:
            result_dict[lll[-10]] = []
            result_dict[lll[-10]].append(lll[-9:])
        else:
            result_dict[lll[-10]].append(lll[-9:])
    return result_dict


def compute_yolo_num(lines_f, name_f, labels_path_f):
    with open(os.path.join(labels_path_f, name_f[:-3] + 'txt')) as fp_f:
        lines1 = fp_f.readlines()
        fp_f.close()
    gt_num_f = len(lines1)
    pred_nums_f = len(lines_f)
    cls_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K']
    img_size = 4096
    pic_true_count = 0
    for line in lines1:
        line = line.split(' ')
        line[-1] = line[-1][:-1]
        line = [eval(kk) for kk in line]
        gt_cls = cls_list[int(line[0])]
        xy4 = xywh2xy4(line, img_size)
        count = 0
        for i in lines_f:
            iou_result = compute_one_ration_iou(xy4, i[1:])
            # print(iou_result)
            if iou_result >= 0.5 and gt_cls == i[0]:
                count = 1
        pic_true_count += count
    with open('/diskb/TZB/TZB1_royolo5/see_per_result.txt', 'a') as p_f:
        sss = '{}:标签数目 {} 预测数目 {} 精确数目 {} 召回率{}\n'.format(name_f, gt_num_f, pred_nums_f, pic_true_count, pic_true_count/gt_num_f)
        p_f.write(sss)
    return pred_nums_f, pic_true_count


def acc():
    pred_result_path = '/diskb/TZB/TZB1_royolo5/inference/results.txt'
    labels_path = '/diskb/TZB/TZB1_royolo5/convertor/fold0/labels/v'

    # 读入预测的结果， 统计预测框的总数量
    with open(pred_result_path) as fp:
        lines0 = fp.readlines()

    pred_dict = pred_lines2dict(lines0)
    pred_count, true_count = 0, 0
    with open('/diskb/TZB/TZB1_royolo5/see_per_result.txt', 'w') as f:
        f.close()
    for pic in pred_dict:
        pred_nums, true_pred_nums = compute_yolo_num(pred_dict[pic], pic, labels_path)
        pred_count += pred_nums
        true_count += true_pred_nums
    # 统计真实标签的总数量
    gt_count = 0
    names_list = os.listdir(labels_path)
    for name in names_list:
        with open(os.path.join(labels_path, name)) as fp:
            lines0 = fp.readlines()
        count0 = len(lines0)
        gt_count += count0
    print("正确预测的个数：", true_count)
    print("总共预测的个数：", pred_count)
    print("真实标签的个数：", gt_count)
    print("召回率： {} %".format((true_count/gt_count) * 100))
    print("精确率： {} %".format((true_count/pred_count) * 100))


def get_gt_dict(labels_path_f):
    cls_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K']
    gt_dict = {'A': 0, 'B': 0, 'C': 0, 'D': 0, 'E': 0, 'F': 0, 'G': 0, 'H': 0, 'I': 0, 'J': 0, 'K': 0}
    names_list_f = os.listdir(labels_path_f)
    for name_f in names_list_f:
        lines_f = open(os.path.join(labels_path_f, name_f)).readlines()
        for line_f in lines_f:
            line_f = line_f.split(' ')
            line_f = [eval(kk) for kk in line_f]
            gt_dict[cls_list[int(line_f[0])]] += 1
    return gt_dict


def get_pred_true_dict(result_dict_f, labels_path_f):
    cls_list = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K']
    pred_true_dict_f = {'A': 0, 'B': 0, 'C': 0, 'D': 0, 'E': 0, 'F': 0, 'G': 0, 'H': 0, 'I': 0, 'J': 0, 'K': 0}
    img_size = 4096
    for pic in result_dict_f:
        with open(os.path.join(labels_path_f, pic[:-3] + 'txt')) as fp:
            lines_f = fp.readlines()
            fp.close()
        for line_f in lines_f:
            line_f = line_f.split(' ')
            line_f = [eval(kk) for kk in line_f]
            gt_cls = cls_list[int(line_f[0])]
            xy4 = xywh2xy4(line_f, img_size)
            for i in result_dict_f[pic]:
                iou_result = compute_one_ration_iou(xy4, i[1:])
                # print(iou_result)
                if iou_result >= 0.5 and gt_cls == i[0]:
                    pred_true_dict_f[gt_cls] += 1
                    break
    return pred_true_dict_f


def get_pred_dict(result_dict_f):
    pred_dict_f = {'A': 0, 'B': 0, 'C': 0, 'D': 0, 'E': 0, 'F': 0, 'G': 0, 'H': 0, 'I': 0, 'J': 0, 'K': 0}
    for name in result_dict_f:
        for one in result_dict_f[name]:
            pred_dict_f[one[0]] += 1
    return pred_dict_f


def dict_count(dict_f):
    count_f = 0
    for ii in dict_f.items():
        count_f += ii[1]
    return count_f


def print_score(dict0, dict1, dict2):
    f_s = 0
    for name_f in dict0:
        if dict1[name_f] == 0:
            p_f = 0
        else:
            p_f = dict2[name_f] / dict1[name_f]
        r_f = dict2[name_f] / dict0[name_f]
        if p_f == 0 and r_f == 0:
            f = 0
        else:
            f = (2*p_f*r_f)/(p_f+r_f)
        f_s += f
        print('类别：{}  P:{}  R:{}  F:{}'.format(name_f, p_f, r_f, f))
    return (f_s / len(dict0)) * 100


def compute_score():
    # pred_result_path = '/data1/TZB/ro_yolov5/test/bicnn/result.txt'
    # pred_result_path = '/data1/TZB/ro_yolov5/test/result.txt'
    pred_result_path = '/data1/TZB/ro_yolov5/test/result_05_500_multi.txt'
    labels_path = '/data1/TZB/ro_yolov5/convertor/org_labels'
    gt_dict = get_gt_dict(labels_path)
    # 读入预测的结果， 统计预测框的总数量
    with open(pred_result_path) as fp:
        lines0 = fp.readlines()
    pred_dict0 = pred_lines2dict(lines0)
    pred_dict = get_pred_dict(pred_dict0)
    pred_true_dict = get_pred_true_dict(pred_dict0, labels_path)
    print('ground_true:', gt_dict)
    print('真实标签的个数：', dict_count(gt_dict))
    print('pred_result:', pred_dict)
    print('总共预测的个数：', dict_count(pred_dict))
    print('pred_true_result:', pred_true_dict)
    print('正确预测的个数：', dict_count(pred_true_dict))
    score = print_score(gt_dict, pred_dict, pred_true_dict)
    print('最后得分：{:.2f}'.format(score))
    # print(dict_count(gt_dict))
    # print(dict_count(pred_true_dict))


# acc()
compute_score()


